fix(geak): forward the run's actual GPU pin in the handoff - #1321

Recipe-first precedence defeats the fix for HIP/CUDA-pinned runs.

materialize_config_with_envs unconditionally synthesizes ROCR_VISIBLE_DEVICES=0..tp-1 into benchmark.envs (_workload_envs.py:943-956), and that materialized recipe is what state.baseline_config_path points at by the time KERNEL runs. So the baseline_recipe source always wins, and it always carries a synthetic 0..tp-1 mask.

Repro: run with HIP_VISIBLE_DEVICES=4,5, TP=2, no ROCR anywhere. _resolve_gpu_pin returns {'var': 'ROCR_VISIBLE_DEVICES', 'value': '0,1', 'ids': [0, 1], 'source': 'baseline_recipe'} (verified by executing the helper). handoff["gpu_ids"] becomes "0,1" where the pre-PR code emitted "4,5", and gpu_pin["value"] = "0,1" tells a ROCR-writing consumer to hard-pin physical cards 0 and 1 — recreating exactly the foreign-tenant card-0 collision #1312 is meant to fix, as a new regression for HIP users.

The recipe mask is only meaningful as a pin when it was authored, not when it was autofilled. Either skip the autofilled ROCR key, or consult the process env first and use the recipe only as a fallback.

The documented {} ("whole machine") case is unreachable in production.

Same root cause as above: because the materialized recipe always carries an autofilled ROCR_VISIBLE_DEVICES=0..tp-1, a genuinely unpinned run still produces a truthy pin. With no mask anywhere in the process env and TP=4, this returns {'ids': [0,1,2,3], 'source': 'baseline_recipe', ...} (verified by executing the helper), so handoff["gpu_pin"] is emitted.

That contradicts the docstring above and the table in docs/components/geak.md ("omitted entirely when no mask is set anywhere, which means whole machine visible, not pinned to card 0"). A GEAK launcher that writes ROCR_VISIBLE_DEVICES = gpu_pin["value"] now restricts a whole-machine run to cards 0-3, where before it inherited everything.

Loop nesting makes source precedence dominate variable precedence.

The source loop is outer and the variable loop inner, so all three vars are checked under baseline_recipe before process_env is reached. That means a recipe-level CUDA_VISIBLE_DEVICES/HIP_VISIBLE_DEVICES silently overrides a real process-level ROCR_VISIBLE_DEVICES pin.

Example: a hand-authored recipe carries benchmark.envs.CUDA_VISIBLE_DEVICES: "0" (a common leftover on a CUDA-derived YAML) while the run is launched with ROCR_VISIBLE_DEVICES=6,7. This returns {'var': 'CUDA_VISIBLE_DEVICES', 'ids': [0]}, and _resolve_handoff_gpu_ids takes the non-ROCR branch and emits gpu_ids="0" verbatim — pinning GEAK to physical card 0 while the run owns 6 and 7.

The docstring's claimed "ROCR_VISIBLE_DEVICES before HIP/CUDA" precedence only holds within a single source. If that ordering is meant to be global, the loops need to be swapped (or the precedence documented as source-major).

A non-numeric but non-blank mask yields a truthy pin with empty ids, and falls back to card 0.

_resolve_gpu_pin returns early on any non-blank string, but _parse_device_list returns [] for anything non-numeric. Two real cases:

• ROCm accepts GPU UUID masks: ROCR_VISIBLE_DEVICES=GPU-a1b2c3,GPU-d4e5f6.
• YAML parses ROCR_VISIBLE_DEVICES: [0,1] as a list, and str() of it is "[0, 1]".

In both, handoff["gpu_pin"] is present with ids: [], this branch sees not ids and returns "0,...,tp-1" — landing every GEAK server on card 0, the exact default the PR exists to eliminate. Meanwhile a consumer reading gpu_pin["ids"] gets nothing, and one writing gpu_pin["value"] exports the literal "[0, 1]".

test_gpu_ids_never_empty_for_a_blank_mask enshrines the fallback as intended, but "never empty" and "never silently card 0" are different goals. Suggest treating a non-blank-but-unparseable mask as an error/warning rather than as unpinned.

The "every layer agrees on the pin" claim does not hold for a present-but-empty mask.

gpu_pool._visible_device_mask returns ([], present=True) for VAR="" and gate.detect_gpu_count() derives 0 GPUs from it. This resolver instead skips a blank value and falls through to the next variable.

Divergence: ROCR_VISIBLE_DEVICES="" together with a stale HIP_VISIBLE_DEVICES="2,3". The orchestrator believes it has zero GPUs, while _resolve_gpu_pin returns {'var': 'HIP_VISIBLE_DEVICES', 'ids': [2, 3]} and hands GEAK two cards. test_pin_skips_blank_values locks the divergent behaviour in.

Either match the other layers (blank = "zero visible", stop) or reword this comment so it doesn't claim an agreement that isn't there.

gpu_ids logical indices are computed against the recipe mask, but the child inherits the process mask.

The GEAK subprocess is launched with runner_env = dict(os.environ) (line ~1171); nothing overrides runner_env["ROCR_VISIBLE_DEVICES"] from the resolved pin. So whenever the recipe mask and the process mask differ, the logical indices resolve against the wrong set.

This is exactly the case the new test_pin_prefers_recipe_over_process_env encodes: recipe ROCR_VISIBLE_DEVICES="6", process ROCR_VISIBLE_DEVICES="0". The handoff advertises gpu_pin.ids=[6] and gpu_ids="0", but the child inherits ROCR=0, so HIP index 0 is physical card 0 — the handoff claims card 6 while every server GEAK launches sits on card 0. If the recipe mask is to be authoritative, the phase should also export it into runner_env.

Separately, tp and gpu_ids can now disagree. handoff["tp"] a few lines up is still int(os.environ.get("TP")) raw, while gpu_ids is clamped to the mask width. With TP=8 exported on a 4-GPU pod the materializer clamps the recipe to TP=4 / ROCR="0,1,2,3" (test_baseline_param_overrides.py:385), so gpu_ids becomes "0,1,2,3" while tp stays 8 — GEAK launches sglang with --tp 8 and four visible cards and fails to load weights. Pre-PR both fields derived from the same $TP and could not disagree; tp should now come from the same resolved mask/recipe as gpu_ids.

These fields are only recorded on the crash-recovery path, so the stated goal is not met.

_geak_reconstruct_from_disk has exactly one call site (line ~906), guarded by if not has_result. A GEAK run that finishes normally and writes geak_result={'status': 'no_gain'} takes the has_result path at line ~900 and returns without ever calling it.

So for every completed run — including the no_gain outcome this comment names — the breakdown contains no gpu_ids/gpu_pin, and a foreign-tenant collision still can't be told apart from a real no_gain. Only a crashed run with no committed result gets the fields. The same two keys need to be recorded on the has_result path too.

Writes an explicit null for unpinned runs and for every v1/v2 handoff.

handoff.get("gpu_pin") yields None for any pre-v3 handoff on disk — still produced by any session resumed from before this deploy — and also for a genuinely unpinned v3 run. A breakdown reader then cannot distinguish "no pin was set" from "this handoff predates the field" from "the pin resolved empty".

This also contradicts the writer's own if gpu_pin: guard in kernel.py and the "omitted when nothing is pinned" contract documented in docs/components/geak.md. Suggest inserting the keys only when present, mirroring the writer.

The end-to-end case never exercises the recipe branch — the branch that wins in production.

This SharedState leaves baseline_config_path unset, so _read_recipe_bench_envs returns {}, _resolve_gpu_pin falls through to the process env, and the source == "process_env" assertion passes.

In every real run baseline_config_path points at a materialized recipe whose benchmark.envs always contains an autofilled ROCR_VISIBLE_DEVICES, so the asserted process_env path is effectively dead code in production — which is why the HIP-pin regression flagged on _resolve_gpu_pin is invisible to the suite.

A case that writes a real materialized YAML (with the autofilled ROCR mask) alongside a HIP_VISIBLE_DEVICES process pin, and asserts the HIP pin survives, would fail today. That is the test this PR most needs.

Two inaccuracies in this row.

1. "capped at tp" is only true for the ROCR branch (min(len(ids), width) in _resolve_handoff_gpu_ids). The HIP/CUDA branch forwards the mask with no cap: HIP_VISIBLE_DEVICES=4,5,6,7 with TP=2 yields gpu_ids="4,5,6,7" — four devices for a two-way tensor-parallel launch.
2. "verbatim" is not accurate either — the ids go through _parse_device_list, which deduplicates and re-serializes. HIP_VISIBLE_DEVICES=" 4, 4 ,5" produces "4,5", not the original string. The _resolve_handoff_gpu_ids docstring makes the same "VERBATIM" claim.

Non-numeric HIP/CUDA mask collapses to 0..tp-1, re-creating the card-0 fallback this PR exists to remove.

With HIP_VISIBLE_DEVICES="GPU-a1b2c3,GPU-d4e5f6" (the ROCm UUID form) and tp=2, _resolve_gpu_pin yields ids=[], count=2. _resolve_handoff_gpu_ids skips the ROCR branch (var != ROCR), falls into if not ids: and returns "0,1". GEAK then exports HIP_VISIBLE_DEVICES=0,1 and every server it launches lands on physical cards 0/1.

Pre-PR the handoff forwarded the UUID string verbatim, so this is a regression on that path. The ROCR branch guards exactly this case with count; the HIP/CUDA branch does not. (Verified by executing the new helpers directly on this branch.)

A present-but-empty mask is treated as "unpinned", so a run with zero visible GPUs is reported as owning the whole machine.

ROCR_VISIBLE_DEVICES="" is read as "zero devices visible" by both policy/gate.detect_gpu_count and bus/gpu_pool._visible_device_mask. Here it is skipped instead: _resolve_gpu_pin returns {}, gpu_pin is omitted from the handoff, and gpu_ids becomes "0,1,..." — GEAK launches on physical card 0, which is the precise collision #1312 is about.

The module comment documents the divergence deliberately, but the chosen fallback is the unsafe one. Emitting a pin with count=0 would at least let the consumer refuse to launch.

Only three mask variables are enumerated, so runs pinned via HSA_VISIBLE_DEVICES or GPU_DEVICE_ORDINAL are still unfixed.

common/env_safety.GPU_MASK_ENV_NAMES lists both alongside the three handled here, and multi_node/scripts/launch_multinode.py:429 strips GPU_DEVICE_ORDINAL as a real mask. A pod pinned with HSA_VISIBLE_DEVICES=7 gets gpu_pin omitted and gpu_ids="0", and every GEAK-launched server lands on physical card 0 — the identical failure, in a change whose whole premise is that a mask variable was overlooked.

_is_autofilled_rocr bails out to False when the recipe carries no TP key, so the synthetic mask still poses as a real pin.

This is the regression the second commit was added to fix, but it only closes when TP is present. Verified on this branch:

_resolve_gpu_pin(recipe_envs={"ROCR_VISIBLE_DEVICES": "0,1"},
                 environ={"HIP_VISIBLE_DEVICES": "4,5"})
-> {var: ROCR, ids: [0, 1], source: baseline_recipe}

Any recipe reaching state.baseline_config_path without a TP key — a hand-authored reference recipe, a _grid_server_args variant where TP came from neither env nor YAML, an on-disk recipe written by an older materializer — re-pins GEAK to cards 0,1 while the run actually owns 4,5.

When both ROCR and HIP are set, the inherited HIP mask is silently discarded and gpu_ids regresses.

ROCR_VISIBLE_DEVICES="4,5,6,7", HIP_VISIBLE_DEVICES="2,3", tp=2: the run actually serves on physical 6,7. Pre-PR gpu_ids="2,3" (correct at the HIP level). Now the ROCR pin wins and gpu_ids="0,1" → physical 4,5, i.e. different cards than the run uses.

cli/preflight._unset_hip_visible_devices only strips the duplicate for CLI-launched runs; the SDK/orchestrator-embedded path, the Ray path and the tests all still reach the KERNEL phase with both set.

Reuse: _VISIBLE_DEVICE_VARS and _parse_device_list are a fifth copy of a tuple/parser pair already in the repo.

• bus/gpu_pool._parse_gpu_list (L60-81) is character-identical to the numeric core of _parse_device_list, and _visible_device_mask (L113) holds the same tuple
• policy/gate.detect_gpu_count (L182) holds it again
• actions/executors/_ray_serving._VISIBLE_DEVICE_ENV_KEYS (L58) is a literal duplicate
• common/env_safety.GPU_MASK_ENV_NAMES (L217) is the dependency-free shared module this belongs in

The in-code comment acknowledges the duplication and defers it, but precedence ordering now has to stay in sync across five sites, and the semantics have already diverged (empty-mask handling, see the comment above).

gpu_ids is clamped to the ROCR mask width but handoff["tp"] is not, so the handoff can ship tp: 2 next to a single device id.

Process ROCR_VISIBLE_DEVICES="6" with TP=2: _workload_envs.materialize_config_with_envs sees len(rocr_devices)=1 < TP=2, expands the recipe to ROCR="0,1" and writes envs["TP"]=2. _resolve_gpu_pin prefers the process env → count=1; _resolve_handoff_gpu_ids returns min(1,2) → "0"; _tp reads the recipe → 2. Handoff is {tp: 2, gpu_ids: "0"} and GEAK launches sglang --tp 2 against one visible card, failing to load weights.

This is the disagreement the adjacent comment states is impossible. The same case also makes gpu_pin report card 6 while the measured baseline actually ran on the recipe's cards 0,1.

gpu_pin["value"] and gpu_ids are only mutually consistent when var == ROCR and source == process_env.

For every other pin, a consumer following the documented contract (write ROCR=value, export HIP=gpu_ids) produces out-of-range HIP indices. With HIP_VISIBLE_DEVICES="4,5", tp=2 → gpu_pin={var: HIP, value: "4,5", ids: [4,5]}, gpu_ids="4,5". adapters/launchers/magpie.sh (named in the PR body as the consumer) writes ROCR_VISIBLE_DEVICES=4,5, so the child sees 2 devices; exporting HIP_VISIBLE_DEVICES=4,5 then indexes positions 4 and 5 of a 2-element set. Identical for a recipe-sourced ROCR pin (absolute gpu_ids="6,7" under ROCR=6,7).

docs/components/geak.md documents the asymmetry, but nothing in the payload lets the consumer detect which case it is holding — source is the only signal, and it is not called out as load-bearing.

The except (TypeError, ValueError) recovery re-runs an unguarded int() and can raise from inside the handler.

Recipe benchmark.envs.TP: "auto" (or 2.0, which int(str(...)) also rejects) plus TP=auto in the process env: the try body raises ValueError, then the handler evaluates int(os.environ.get("TP", "1") or 1) and raises again during handling, propagating out of _run_geak_kernel_phase and killing the KERNEL phase.

The fallback exists precisely because the primary parse is untrusted, so it should not repeat the same untrusted parse bare.

on_error fires on a missing geak/handoff.json, adding a spurious warning to healthy runs.

FileNotFoundError is an OSError, and common/jsonio.read_json invokes on_error(exc) for OSError before returning the default. Any resumed pre-v3 session, or a run whose GEAK artifacts live under an external exp_root, reaches this unconditional read on the SUCCESS collect path and gets geak: handoff read failed: [Errno 2] No such file or directory appended to the breakdown warnings.

It is also unconditional disk I/O on a path that previously touched no files, for two optional fields — worth gating on the file existing, or passing an on_error that ignores FileNotFoundError.